import os
import matplotlib.pyplot as plt


class KneeDataset:
    def __init__(self, image_dir, label_dir=None):
        if label_dir is None:
            # 处理只有图像目录的情况
            self.images = [os.path.join(image_dir, img) for img in os.listdir(image_dir)]
        else:
            # 处理有图像和标签目录的情况
            self.images = [os.path.join(image_dir, img) for img in os.listdir(image_dir)]
            self.labels = [os.path.join(label_dir, lbl) for lbl in os.listdir(label_dir)]


def merge_and_save_images(image_dir, label_dir, prediction_dir, output_dir):
    """
    合并原图像、标签图像和预测图像，并保存到指定文件夹。

    参数:
    image_dir (str): 原图像所在目录。
    label_dir (str): 标签图像所在目录。
    prediction_dir (str): 预测图像所在目录。
    output_dir (str): 合并图像保存的目录。
    """
    # 创建输出文件夹
    os.makedirs(output_dir, exist_ok=True)

    # 加载数据集（这里直接获取图像和标签数据）
    image_dataset = KneeDataset(image_dir)
    label_dataset = KneeDataset(label_dir)
    prediction_dataset = KneeDataset(prediction_dir, None)

    for idx, (image, label, prediction) in enumerate(zip(image_dataset.images, label_dataset.images, prediction_dataset.images)):
        # 创建子图布局
        fig, axes = plt.subplots(1, 3, figsize=(15, 5))
        # 显示原图像
        axes[0].imshow(plt.imread(image), cmap='gray')
        axes[0].set_title('Original Image')
        axes[0].axis('off')
        # 显示标签图像
        axes[1].imshow(plt.imread(label), cmap='gray')
        axes[1].set_title('Ground Truth')
        axes[1].axis('off')
        # 显示预测图像
        axes[2].imshow(plt.imread(prediction), cmap='gray')
        axes[2].set_title('Prediction')
        axes[2].axis('off')
        # 保存合并后的图像
        save_path = os.path.join(output_dir, f"merged_{idx}.png")
        plt.savefig(save_path)
        plt.close(fig)


if __name__ == "__main__":
    image_dir = "test_image"
    label_dir = "test_label"
    prediction_dir = "test_results"
    output_dir = "merged_images"
    merge_and_save_images(image_dir, label_dir, prediction_dir, output_dir)